//------------------------------------------------------------------------------
// <copyright file="RpcServerImpl.cpp" company="Microsoft">
//     Copyright (c) Microsoft Corporation.  All rights reserved.
// </copyright>
//------------------------------------------------------------------------------

#include "stdafx.h"

#include "RpcServer.h"
#include "RpcServerImpl.h"
#include "RpcErrorCodes.h"
#include "RpcArchive.h"
#include "RpcArchiveImpl.h"
#include "RpcRequest.h"
#include "RpcResponse.h"
#include "RpcMessage.h"
#include "IRpcChannel.h"
#include "RpcErrorCodes.h"
#include "RpcMessageDispatcher.h"
#include "RpcBuiltInHooks.h"
#include <sstream>

using namespace std;
using namespace Microsoft::Internal::GamesTest::Rpc::Server;

// {21A25F93-2841-4CC7-A380-FD91ABBE43D1}
const GUID RpcServerImpl::TdkRpcDefaultExtensionId = { 0x21a25f93, 0x2841, 0x4cc7, { 0xa3, 0x80, 0xfd, 0x91, 0xab, 0xbe, 0x43, 0xd1 } };

// Default listening port
const int DefaultRpcListeningPort = 4600;

// Max. queue size
//
// We determined the average message size to be about 372 bytes:
//   338 (serialized from client) + 32 (PendingIncomingRequest overhead) + 2 (concurrent queue, 16 bytes per page of 8)
// We know that the concurrent_queue stores PendingIncomingRequest objects in groups of 8, so we want the maximum
// number of requests in the queue to be a multiple of 8.  We'd also like to keep the queue size around 512 kilobytes.
// 512Kb / 372 bytes = 1409.376
// 1408 is the nearest number that is also a multiple of 8.
//
const unsigned int DefaultRpcMaxQueueSize = 1408;

// Static member used to disable RPC on production sandboxes.
RpcServerImpl::EnabledStateType RpcServerImpl::s_enabledState = RpcServerImpl::EnabledStateType::Uninitialized;

RpcServerImpl::RpcServerImpl()
	: m_extensionId(TdkRpcDefaultExtensionId),
	  m_numConnectedClients(0),
	  m_isStarted(false),
	  m_port(DefaultRpcListeningPort),
	  m_maxQueueSize(DefaultRpcMaxQueueSize),
	  m_immediateHooksRegistered(false)
{
}

RpcServerImpl::RpcServerImpl(int port)
	: m_extensionId(TdkRpcDefaultExtensionId),
	  m_numConnectedClients(0),
	  m_isStarted(false),
	  m_port(port),
	  m_maxQueueSize(DefaultRpcMaxQueueSize),
	  m_immediateHooksRegistered(false)
{
}

RpcServerImpl::RpcServerImpl(const GUID& serverId)
	: m_extensionId(serverId),
	  m_numConnectedClients(0), 
	  m_isStarted(false),
	  m_port(DefaultRpcListeningPort),
	  m_maxQueueSize(DefaultRpcMaxQueueSize),
	  m_immediateHooksRegistered(false)
{
}

RpcServerImpl::RpcServerImpl(const GUID& serverId, int port)
	: m_extensionId(serverId),
	  m_numConnectedClients(0),
	  m_isStarted(false),
	  m_port(port),
	  m_maxQueueSize(DefaultRpcMaxQueueSize),
	  m_immediateHooksRegistered(false)
{
}

RpcServerImpl::RpcServerImpl(unsigned int maxQueueSize)
	: m_extensionId(TdkRpcDefaultExtensionId),
	  m_numConnectedClients(0),
	  m_isStarted(false),
	  m_port(DefaultRpcListeningPort),
	  m_maxQueueSize(maxQueueSize),
	  m_immediateHooksRegistered(false)
{
}

RpcServerImpl::RpcServerImpl(int port, unsigned int maxQueueSize)
	: m_extensionId(TdkRpcDefaultExtensionId),
	  m_numConnectedClients(0),
	  m_isStarted(false),
	  m_port(port),
	  m_maxQueueSize(maxQueueSize),
	  m_immediateHooksRegistered(false)
{
}

RpcServerImpl::RpcServerImpl(const GUID& serverId, unsigned int maxQueueSize)
	: m_extensionId(serverId),
	  m_numConnectedClients(0), 
	  m_isStarted(false),
	  m_port(DefaultRpcListeningPort),
	  m_maxQueueSize(maxQueueSize),
	  m_immediateHooksRegistered(false)
{
}

RpcServerImpl::RpcServerImpl(const GUID& serverId, int port, unsigned int maxQueueSize)
	: m_extensionId(serverId),
	  m_numConnectedClients(0),
	  m_isStarted(false),
	  m_port(port),
	  m_maxQueueSize(maxQueueSize),
	  m_immediateHooksRegistered(false)
{
}

HRESULT RpcServerImpl::ProcessPendingRequests()
{
	if(!GetIsEnabled())
	{
		return GAMESTEST_RPC_E_METHOD_DISABLED_INVALID_SANDBOX;
	}

	std::unique_lock<std::mutex> ul(m_reEntryLock, std::try_to_lock);

	if (!ul.owns_lock())
	{
		return S_OK;
	}

	PendingOutgoingRequest pendingOutgoingRequest;
	PendingIncomingRequest pendingIncomingRequest;

	bool anyRequestSucceeded = false;
	bool anyRequestFailed = false;

	while(m_pendingOutgoingRequests.try_pop(pendingOutgoingRequest))
	{
		HRESULT hr = ProcessPendingRequest(pendingOutgoingRequest);

		if(FAILED(hr))
		{
			RpcPrintMessage(L"Failed to queue outgoing request. Remote method name: ");
			RpcPrintMessage(pendingOutgoingRequest.eventName.c_str());

			RpcPrintError(hr);

			anyRequestFailed = true;
		}
		else
		{
			anyRequestSucceeded = true;
		}
	}

	while(m_pendingIncomingRequests.try_pop(pendingIncomingRequest))
	{
		HRESULT hr = ProcessPendingRequest(pendingIncomingRequest);

		if(FAILED(hr))
		{
			RpcPrintMessage(L"Failed to fully respond to incoming request. ");
			if(!pendingIncomingRequest.request.empty())
			{
				RpcPrintMessage(L"Request Content: ");
				RpcPrintMessage(pendingIncomingRequest.request.c_str());
			}

			RpcPrintError(hr);

			anyRequestFailed = true;
		}
		else
		{
			anyRequestSucceeded = true;
		}
	}

	if(!anyRequestSucceeded && !anyRequestFailed)
	{
		// No requests were processed
		return S_OK;
	}
	else if(!anyRequestSucceeded && anyRequestFailed)
	{
		// Every request failed
		return GAMESTEST_RPC_E_PROCESS_PENDING_REQUEST_FAILURE;
	}
	else if(anyRequestSucceeded && anyRequestFailed)
	{
		// Some requests succeeded and some requests failed
		return GAMESTEST_RPC_I_PROCESS_PENDING_REQUEST_PARTIAL_FAILURE;
	}
	else
	{
		_ASSERTE(anyRequestSucceeded && !anyRequestFailed && "All requests should have succeeded.");
		// Every request succeeded
		return S_OK;
	}
}

HRESULT RpcServerImpl::ProcessPendingRequest(const PendingIncomingRequest& pendingRequest)
{
	HRESULT hr = S_OK;

	auto channel = pendingRequest.channel.lock();

	if (channel != nullptr)
	{
		RpcRequest deserializedRequest;
		RpcResponse response;
		RpcMessage responseMessage;
		std::wstring serializedResponse;
		std::wstring serializedMessage;
		RpcHookTable::iterator hookIter = m_hookTable.end();

		if (pendingRequest.request.empty())
		{
			hr = E_INVALID_PROTOCOL_FORMAT;
			goto end;
		}

		CHK(RpcRequest::Parse(pendingRequest.request, &deserializedRequest));

		response.RequestId = deserializedRequest.RequestId;

		hookIter = m_hookTable.find(deserializedRequest.MethodName);

		if(hookIter == m_hookTable.end())
		{
			response.RpcErrorCode = GAMESTEST_RPC_E_METHOD_NOT_FOUND;
			goto end;
		}

#pragma warning( push )
#pragma warning( disable : 4239 )
		response.HookErrorCode = hookIter->second.first(deserializedRequest.Parameters, response.Archive, channel);
#pragma warning ( pop )

end:
		if(SUCCEEDED(hr))
		{
			// We can't use the "CHK" macro like normal once we are inside of the "end"
			// section because if something returns a failure code, then we'll jump back to
			// the "end" label and then run the same failing code again.  This will likely
			// lead to an infinite loop.
			hr = response.Serialize(serializedResponse);

#pragma warning(suppress: 6221)
			if(SUCCEEDED(hr) || hr == GAMESTEST_RPC_W_FAILED_TO_SERIALIZE_USER_DATA)
			{
				responseMessage.messageType = MessageType::Response;
				responseMessage.payload = std::move(serializedResponse);
				hr = responseMessage.Serialize(serializedMessage);

#pragma warning(suppress: 6221)
				if (SUCCEEDED(hr) || hr == GAMESTEST_RPC_W_FAILED_TO_SERIALIZE_USER_DATA)
				{
					hr = channel->SendRpcMessage(serializedMessage, this->ServerId, hr);
				}
			}
		}

		// This is not an "else" condition. If any of the above calls within the "end" section fail, we want to send the error message.
#pragma warning(suppress: 6221)
		if(FAILED(hr) && hr != GAMESTEST_RPC_W_FAILED_TO_SERIALIZE_USER_DATA)
		{
			hr = channel->SendRpcMessage(L"", this->ServerId, hr);
		}
	}
	else
	{
		hr = GAMESTEST_RPC_E_CLIENT_DISCONNECTED;
	}

	return hr;
}

HRESULT RpcServerImpl::ProcessPendingRequest(PendingOutgoingRequest& pendingRequest)
{
	HRESULT hr = S_OK;

	std::wstring serializedRequest;
	std::wstring serializedMessage;
	bool anyRequestSuccess = false;

	std::pair<RpcSubscriberTable::iterator, RpcSubscriberTable::iterator> range =
		m_subscriberTable.equal_range(pendingRequest.eventName);

	if (range.first != m_subscriberTable.end())
	{
		RpcRequest rpcRequest(std::move(pendingRequest.eventName), std::move(pendingRequest.archive));
		CHK(rpcRequest.Serialize(serializedRequest));

		RpcMessage rpcMessage;
		rpcMessage.messageType = MessageType::Request;
		rpcMessage.payload = std::move(serializedRequest);
		CHK(rpcMessage.Serialize(serializedMessage));

		for(RpcSubscriberTable::iterator iter = range.first; iter != range.second; /*iter only incremented if element isn't erased. */)
		{
			std::shared_ptr<IRpcChannel> channel = (*iter).second.lock();
			if (channel != nullptr)
			{
				HRESULT sendHr = channel->SendRpcMessage(serializedMessage, ServerId, 0);

				if (FAILED(sendHr))
				{
					RpcPrintMessage(L"Failed to raise event to a client. ");
					RpcPrintError(hr);

					hr = GAMESTEST_RPC_I_PROCESS_PENDING_REQUEST_PARTIAL_FAILURE;
				}
				else
				{
					anyRequestSuccess = true;
				}

				iter++;
			}
			else
			{
				iter = m_subscriberTable.erase(iter);
			}
		}
	}
end:
#pragma warning(suppress: 6221)
	if (hr == GAMESTEST_RPC_I_PROCESS_PENDING_REQUEST_PARTIAL_FAILURE && !anyRequestSuccess)
	{
		return GAMESTEST_RPC_E_PROCESS_PENDING_REQUEST_FAILURE;
	}

	return hr;
}

HRESULT RpcServerImpl::ProcessIncomingMessage(std::weak_ptr<IRpcChannel> channel, const std::wstring& message)
{
	HRESULT hr = S_OK;
	RpcMessage rpcMessage;

	CHK(RpcMessage::Parse(message, &rpcMessage));

	if (rpcMessage.messageType == MessageType::Request)
	{
		CHK(QueueIncomingRequest(channel, std::move(rpcMessage.payload)));
	}

	// Responses are currently unused by Server side.
	// These are responses to events raised, which do not require a response.
end:
	return hr;
}

void RpcServerImpl::ClientDisconnected(std::weak_ptr<IRpcChannel> client)
{
	// We don't want to unsubscribe while we're in the middle of processing requests.
	std::unique_lock<std::mutex> ul(m_reEntryLock);

	std::shared_ptr<IRpcChannel> sharedClient = client.lock();
	if(sharedClient != nullptr)
	{
		for (auto iter = m_subscriberTable.begin(); iter != m_subscriberTable.end(); /* iter only incremented if element is not removed. */ )
		{
			auto clientIter = iter->second.lock();
			if (clientIter == nullptr || clientIter.get() == sharedClient.get())
			{
				iter = m_subscriberTable.erase(iter);
			}
			else
			{
				iter++;
			}
		}
	}
}

HRESULT RpcServerImpl::QueueIncomingRequest(std::weak_ptr<IRpcChannel> channel, std::wstring&& request)
{
	HRESULT hr = S_OK;

	if (m_immediateHooksRegistered)
	{
		// Find out if the request is for a method that has to be processed immediately on the receiving thread.
		std::wstring methodName;
		CHK(RpcRequest::PreParse(request, methodName));

		RpcHookTable::iterator hookIter = m_hookTable.find(methodName);

		if (hookIter == m_hookTable.end())
		{
			hr = GAMESTEST_RPC_E_METHOD_NOT_FOUND;
			goto end;
		}

		if (hookIter->second.second)
		{
			// Process request directly on the current thread (receiving thread for this client).
			PendingIncomingRequest immediateRequest;
			immediateRequest.channel = channel;
			immediateRequest.request = std::move(request);

			hr = ProcessPendingRequest(immediateRequest);
			goto end;
		}
	}

	size_t currentIncomingQueueSize = m_pendingIncomingRequests.unsafe_size();
	size_t currentOutgoingQueueSize = m_pendingOutgoingRequests.unsafe_size();

	if((currentIncomingQueueSize + currentOutgoingQueueSize) < m_maxQueueSize)
	{
		PendingIncomingRequest pendingRequest;
		pendingRequest.channel = channel;
		pendingRequest.request = std::move(request);

		m_pendingIncomingRequests.push(pendingRequest);
	}
	else
	{
		hr = GAMESTEST_RPC_E_PENDING_REQUESTS_QUEUE_FULL;
	}

end:
	return hr;
}

GAMESTEST_ANNOTATE_HRESULT_METHOD HRESULT RpcServerImpl::RaiseEvent(_In_z_ const wchar_t* eventName, _In_ RpcArchiveImplPtr archive)
{
	if(!GetIsEnabled())
	{
		return GAMESTEST_RPC_E_METHOD_DISABLED_INVALID_SANDBOX;
	}

	HRESULT hr = E_FAIL;
	VERIFY_PARAM_NOT_NULL(eventName);

	size_t currentIncomingQueueSize = m_pendingIncomingRequests.unsafe_size();
	size_t currentOutgoingQueueSize = m_pendingOutgoingRequests.unsafe_size();

	if((currentIncomingQueueSize + currentOutgoingQueueSize) < m_maxQueueSize)
	{
		m_pendingOutgoingRequests.push(PendingOutgoingRequest(eventName, std::move(archive)));
		hr = S_OK;
	}
	else
	{
		hr = GAMESTEST_RPC_E_PENDING_REQUESTS_QUEUE_FULL;
	}

end:
	return hr;
}

GAMESTEST_ANNOTATE_HRESULT_METHOD HRESULT RpcServerImpl::Start()
{
	if(!GetIsEnabled())
	{
		return GAMESTEST_RPC_E_METHOD_DISABLED_INVALID_SANDBOX;
	}

	HRESULT hr;
	if(m_isStarted)
	{
		return GAMESTEST_RPC_S_ALREADY_LISTENING;
	}
	else
	{
		CHK(RpcMessageDispatcher::RegisterMessageHandler(this->ServerId, this->Port, this));
	}

	CHK(RegisterBuiltInHookMethods());

end:
	if(FAILED(hr))
	{
		m_isStarted = false;
	}
	else
	{
		m_isStarted = true;
	}

	return hr;
}

GAMESTEST_ANNOTATE_HRESULT_METHOD HRESULT RpcServerImpl::Stop()
{
	if(!GetIsEnabled())
	{
		return GAMESTEST_RPC_E_METHOD_DISABLED_INVALID_SANDBOX;
	}

	HRESULT hr;

	if(!m_isStarted)
	{
		return GAMESTEST_RPC_S_NOT_STARTED;
	}
	else
	{
		CHK(RpcMessageDispatcher::UnregisterMessageHandler(this->ServerId));
	}

end:
	if(SUCCEEDED(hr))
	{
		m_isStarted = false;
	}

	return hr;
}

bool RpcServerImpl::GetIsStarted()
{
	return m_isStarted;
}

const GUID& RpcServerImpl::GetServerId()
{
	return m_extensionId;
}

const GUID& RpcServerImpl::GetDefaultServerId()
{
	return TdkRpcDefaultExtensionId;
}

int RpcServerImpl::GetPort()
{
	return m_port;
}

int RpcServerImpl::GetDefaultPort()
{
	return DefaultRpcListeningPort;
}

bool RpcServerImpl::GetIsEnabled()	
{
#ifdef _DURANGO
	// s_enabledState is a enum so that we can delay initialization until the first time someone asks for it
	// which realistically speaking would probably be a call to RegisterHookMethod. We are delaying so that
	// any error which may occur attempting to get the sandbox id will occur when the user is calling a method,
	// not in some static initialization phase which could cause issues that would be hard to track.
	switch(s_enabledState)
	{
	case EnabledStateType::Uninitialized:
		{
			const std::wstring RetailSandboxid = L"RETAIL";
			const std::wstring CertSandboxId = L"CERT";
			const std::wstring CertSandboxSuffix = L".99";
			const int CertSandboxSuffixLength = 3;

			std::wstring sandbox = GetSandboxId();

			if(sandbox.length() > 3)
			{
				if(RetailSandboxid == sandbox.data() 
					|| CertSandboxId == sandbox.data() 
					|| CertSandboxSuffix == (sandbox.data() + sandbox.length() - CertSandboxSuffixLength))
				{
					s_enabledState = EnabledStateType::Disabled;
					return false;
				}
			}
			s_enabledState = EnabledStateType::Enabled;
			return true;
		}
		break;
	case EnabledStateType::Disabled:
		return false;
	case EnabledStateType::Enabled:
		return true;
	}
#endif
	return true;
}

/// <summary>
/// Gets the default max. queue size for RpcServer.
/// </summary>
unsigned int RpcServerImpl::GetDefaultMaxQueueSize()
{
	return DefaultRpcMaxQueueSize;
}

#ifdef _DURANGO
std::wstring RpcServerImpl::GetSandboxId()
{
	return std::wstring(Windows::Xbox::Services::XboxLiveConfiguration::SandboxId->Data());
}
#endif

GAMESTEST_ANNOTATE_HRESULT_METHOD HRESULT RpcServerImpl::RegisterHookMethod(_In_z_ const wchar_t* hookName, _In_ RpcHookFunction hookFunc, bool immediate)
{
	if(!GetIsEnabled())
	{
		return GAMESTEST_RPC_E_METHOD_DISABLED_INVALID_SANDBOX;
	}

	HRESULT hr = S_OK;

	VERIFY_PARAM_NOT_NULL(hookName);

	if (immediate)
	{
		// Verify hook is not already registered
		RpcHookTable::iterator exists = m_hookTable.find(hookName);

		if (exists != m_hookTable.end() && exists->second.second != true)
		{
			hr = GAMESTEST_RPC_E_IMMEDIATE_METHOD_ALREADY_REGISTERED;
			goto end;
		}

		m_immediateHooksRegistered = true;
	}

	m_hookTable[hookName] = std::pair<RpcHookFunction, bool>(hookFunc, immediate);

end:
	return hr;
}

HRESULT RpcServerImpl::RegisterBuiltInHookMethods()
{
	HRESULT hr = E_FAIL;

	CHK(RegisterHookMethod(L"_52DA7DDB-51D8-4DB7-B447-02F8CCB4DF81_VerifyPropertiesAndSubscribeOnConnect",
		[=](RpcArchiveImpl& input, RpcArchiveImpl& output, std::weak_ptr<IRpcChannel> clientChannel) -> RpcHookCode
	{
		return VerifyPropertiesAndSubscribeOnConnect(input, output, clientChannel, this);
	}, false));

	CHK(RegisterHookMethod(L"_3FA5BC26-5FE9-40D5-9080-330582E93449_SetMaximumQueueSize",
		[=](RpcArchiveImpl& input, RpcArchiveImpl& output, std::weak_ptr<IRpcChannel> clientChannel) -> RpcHookCode
	{
		UNREFERENCED_PARAMETER(clientChannel);
		return SetMaximumQueueSize(input, output, this);
	}, false));

	CHK(RegisterHookMethod(L"_2EB24267-D833-4D0E-879F-3637D0C909A1_UpdateEventSubscription",
		[=](RpcArchiveImpl& input, RpcArchiveImpl& output, std::weak_ptr<IRpcChannel> clientChannel) -> RpcHookCode
	{
		return UpdateEventSubscription(input, output, clientChannel, this);
	}, false));

	CHK(RegisterHookMethod(L"_7DE6D53C-7AE7-49C2-9FC5-95F93CF2028E_EnumRegisteredHooks",
		[=](RpcArchiveImpl& input, RpcArchiveImpl& output, std::weak_ptr<IRpcChannel> clientChannel) -> RpcHookCode
	{
		return EnumRegisteredHooks(input, output, clientChannel, this);
	}, false));

end:
	return hr;
}

#ifdef _WINRT_DLL

RpcServerImpl* RpcServerImpl::GetImplementation(RpcServer^ server)
{
	return server->m_impl;
}

#else

RpcServerImpl* RpcServerImpl::GetImplementation(RpcServer* server)
{
	return server->m_impl;
}

#endif

RpcServerImpl::~RpcServerImpl()
{
	if(m_isStarted)
	{
#pragma warning(suppress: 6031)
		RpcMessageDispatcher::UnregisterMessageHandler(this->ServerId);
	}
}